package buaa.qel.rpc.server.registry;

import buaa.qel.rpc.common.info.RegistryNodeInfo;
import buaa.qel.rpc.common.info.RpcServiceInfo;
import buaa.qel.rpc.common.util.ServiceKeyMaker;
import buaa.qel.rpc.common.zookeeper.CuratorClient;
import buaa.qel.rpc.common.zookeeper.config.ZookeeperConstants;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.state.ConnectionState;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * ServiceRegistry
 */
@Slf4j
public class ServiceRegistry {
    private CuratorClient curatorClient;
    private List<String> pathList = new ArrayList<>();

    public ServiceRegistry(String registryAddress) {
        this.curatorClient = new CuratorClient(registryAddress, 5000);
    }

    public void registerService(String host, int port, Map<String, Object> serviceMap) {
        // Register service info
        List<RpcServiceInfo> serviceInfoList = new ArrayList<>();
        for (String key : serviceMap.keySet()) {
            String[] serviceInfo = key.split(ServiceKeyMaker.SERVICE_KEY_SALT);
            if (serviceInfo.length > 0) {
                RpcServiceInfo rpcServiceInfo = new RpcServiceInfo();
                rpcServiceInfo.setInterfaceName(serviceInfo[0]);
                if (serviceInfo.length == 2) {
                    rpcServiceInfo.setServiceName(serviceInfo[1]);
                } else {
                    rpcServiceInfo.setServiceName("");
                }
                log.info("Register new service: {} ", key);
                serviceInfoList.add(rpcServiceInfo);
            } else {
                log.warn("Can not get service name and version: {} ", key);
            }
        }
        try {
            log.info("Try to create zookeeper node...");
            RegistryNodeInfo registerInfo = new RegistryNodeInfo();
            registerInfo.setHost(host);
            registerInfo.setPort(port);
            registerInfo.setServiceInfoList(serviceInfoList);
            String serviceData = registerInfo.toJson();
            byte[] bytes = serviceData.getBytes();
            String path = ZookeeperConstants.ZK_DATA_PATH + "-" + registerInfo.hashCode();
            log.info("Try to create zookeeper node {}",path);
            path = this.curatorClient.createPathData(path, bytes);
            pathList.add(path);
            log.info("Register {} new service, host: {}, port: {}", serviceInfoList.size(), host, port);
        } catch (Exception e) {
            log.error("Register service fail, exception: {}", e.getMessage());
        }

        curatorClient.addConnectionStateListener((CuratorFramework curatorFramework, ConnectionState connectionState) -> {
                if (connectionState == ConnectionState.RECONNECTED) {
                    log.info("Connection state: {}, register service after reconnected", connectionState);
                    registerService(host, port, serviceMap);
                }
            });
    }

    public void unregisterService() {
        log.info("Unregister all service");
        for (String path : pathList) {
            try {
                this.curatorClient.deletePath(path);
            } catch (Exception ex) {
                log.error("Delete service path error: " + ex.getMessage());
            }
        }
        this.curatorClient.close();
    }
}
